package util

import (
	"errors"
	"fmt"
	"slices"
	"sort"
	"sync"
)

type TreeItem struct {
	Data      any         `json:"data,omitempty"`
	Label     string      `json:"label"`
	Key       string      `json:"key"`
	ParentKey string      `json:"parentKey,omitempty"`
	Children  []*TreeItem `json:"children,omitempty"`
	Rank      uint        `json:"rank,omitempty"`
	Level     uint8       `json:"level,omitempty"`
}

type Tree struct {
	items map[string]*TreeItem
	roots []*TreeItem
	mutex sync.Mutex
}

var ErrTreeParentNotFound = errors.New("tree parent not found")

func NewTreeItem(key, label string, rank uint, data any) (*TreeItem, error) {
	if key == "" || key == "-" {
		return nil, fmt.Errorf("key is empty")
	}
	if label == "" {
		label = key
	}
	return &TreeItem{
		Data:     data,
		Label:    label,
		Key:      key,
		Children: make([]*TreeItem, 0),
		Rank:     rank,
	}, nil
}

func NewTree() *Tree {
	return &Tree{
		items: make(map[string]*TreeItem),
		roots: make([]*TreeItem, 0),
	}
}

func (t *Tree) AddItem(item *TreeItem, parentKey string) error {
	t.mutex.Lock()
	defer t.mutex.Unlock()
	key := item.Key
	if _, ok := t.ItemByKey(key); ok {
		return fmt.Errorf("item with key %s already exists", key)
	}
	if parentKey != "" {
		pItem, ok := t.items[parentKey]
		if !ok || pItem == nil {
			return ErrTreeParentNotFound
		}
		item.Level = pItem.Level + 1
		item.ParentKey = parentKey
		pItem.Children = append(pItem.Children, item)
		sort.Slice(pItem.Children, func(i, j int) bool {
			return pItem.Children[i].Rank > pItem.Children[j].Rank
		})
	} else {
		t.roots = append(t.roots, item)
		sort.Slice(t.roots, func(i, j int) bool {
			return t.roots[i].Rank > t.roots[j].Rank
		})
	}
	t.items[key] = item
	return nil
}

func (t *Tree) Export() []*TreeItem {
	return t.roots
}

// ExportForSearch 空表示“全部”，减号表示“未归类”
func (t *Tree) ExportForSearch() []*TreeItem {
	ret := &TreeItem{
		Label:    "- 全部 -",
		Key:      "",
		Children: make([]*TreeItem, 0, len(t.roots)+1),
	}
	ret.Children = append(ret.Children, &TreeItem{
		Label: "- 未归类 -",
		Key:   "-",
	})
	ret.Children = append(ret.Children, t.roots...)
	res := make([]*TreeItem, 0, 1)
	res = append(res, ret)
	return res
}

func (t *Tree) ItemByKey(key string) (*TreeItem, bool) {
	item, ok := t.items[key]
	return item, ok
}

func (t *Tree) ChildrenKeys(parentKey string) []string {
	ret := make([]string, 0)
	item, ok := t.ItemByKey(parentKey)
	if !ok {
		return ret
	}
	for _, childItem := range item.Children {
		ret = append(ret, childItem.Key)
		ret = append(ret, t.ChildrenKeys(childItem.Key)...)
	}
	return ret
}

func (t *Tree) _ancestors(key string) []*TreeItem {
	ret := make([]*TreeItem, 0)
	item, ok := t.ItemByKey(key)
	if !ok || item.ParentKey == "" {
		return ret
	}
	if itemParent, ok := t.ItemByKey(item.ParentKey); ok {
		ret = append(ret, itemParent)
		ret = append(ret, t._ancestors(itemParent.Key)...)
	}
	return ret
}

func (t *Tree) Ancestors(key string) []*TreeItem {
	ret := t._ancestors(key)
	if len(ret) > 1 {
		slices.Reverse(ret)
	}
	return ret
}
